/*
 * This file is part of the openHiTLS project.
 *
 * openHiTLS is licensed under the Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *
 *     http://license.coscl.org.cn/MulanPSL2
 *
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
 */
#include "hitls_build.h"
#if defined(HITLS_CRYPTO_CMVP_ISO19790) || defined(HITLS_CRYPTO_CMVP_FIPS)
#include <string.h>
#include "crypt_cmvp_selftest.h"
#include "cmvp_common.h"
#include "err.h"
#include "crypt_errno.h"
#include "crypt_eal_pkey.h"
#include "crypt_eal_rand.h"
#include "bsl_err_internal.h"
#include "crypt_utils.h"
#include "crypt_util_rand.h"
#include "securec.h"
#include "bsl_sal.h"

#define BITS_OF_BYTE 8

typedef struct {
    const char *ek;
    const char *dk;
    const char *c;
    const char *k;
    const char *m;
    int32_t algId;
    int32_t type;
} CMVP_MlkemVector;

static const CMVP_MlkemVector MLKEM_VECTOR[] = {
    {
        // https://github.com/usnistgov/ACVP-Server/tree/master/gen-val/json-files/ML-KEM-encapDecap-FIPS203
        .ek = "EDDBA9CB830474EC0111408243A97335494DDDFA8D5C6A344B0C9ED1061A7D4A840891C24E140BE6D7BDEE01571A1312D4000591F2224B850105878FD3966F537AC348D042B95785B8171F18647619072780C131BDAA462D34A039014012BBB2197A4328F89CC5B750CAD9495BE607D5017DFA105D87C830432795B42C5E89D8B536C087592829ED934822768A1824696A278952627CA2EA73CDA4A306F972B8521B6A7B28DD0AC44E6A8C46BA7291D48574F8356065AFBB073DF9B84B520673F6644D3959BAD576497E0BC846F560CE697C35493A995597F66315639A7B20B84105EAA48F3B4593A0BF12D64AD4CB043B6435D44A2F5E136F21621B88704DF8855C4F1689761A239DFA6874455D48251EB087865AC3825C275A1018B112F01E82D8698F684C8EE127D9D25FB92B180677A36930793DAB0306AAC8BC80A79885B0737C6B319600BB859966AB9F082490E31B669008CD055921CB2360667180611485E5D70E774401739B16F54601D4A600608119E4B3112AAC92D1C094E40318D4259044D16D3AD62B35170B54D16F0F850C593B2062CC3A769B4D28B6A776F981D4BC835612349292B748E30098D29C668A80F68718BF40B7E785CE41C094A7326F3E6992D99B80E64350C5395823B1315160582EC10B01A4337C52C355E98B857B2DF0006A0E8AB49F5BA23B8C7E329C20F91386852A8948751C0F301465D713E7857A7BD8B6076C059BE378F74BB257750F544A31874B79ABEA78F8647E77001506919AA467019E924B7BE53705A5C9411A8B513C8C51300849136E77919CF83CBA9EA1CD6004AA7178C8E7D062E5E13544E060B34984F61C2EB6B56580C45F67DA8982A97BD5BBA9D98462161A55357079C15A2A33992D4B66093DB03676D0C385DCC6DD3C18EC8138738AA0FD6058C9F721C5E3131D7828C9675048949D0CE32E92A3363EC0527F54C89B5140590822D9C96CB1BAA43A8CACCC085EA6FB602B871128AB2C8A2A1644623A1D7175CF4A23F4B497B7BB4D0B652EA4272D7A99A9A5F4C54EAA6482D0BFB4E46EE482B892895C78DCBAE4414543E0BF8AC4C82889BC6DDC385624AC7805CCC644BB430588AA19723AEE5999EF2AE2565C3BB3",
        .dk = "502B723FE238EDE6A7867C66B49BAA99280584DB1318281150128D65F8A2D54381444994AD938805A3065271BCB390B9D100454A5644C5E858C1F223BF83654627AB98F433AED20E26D8597CD41D802C6323B18446487E51D3660458156DF1571E092763133B87F35BDC2CBC39422A65FB36488B523123B36E3AAF7A8C4DFFAA3719460ACEEC6F5F6957619B776925854A7A49E976592B042A7D3C49F3089F3FAC6F29185BF5DCAF77E6AB69F7ABB926A360385F5A79B7CB8CCC63BC38E0495A89A70E461CBDDBE34AA2F6CE7BF3C1533569B61A88DC1038A5B27C42D024EAF60041D0C771833D3C3213572CBC1C356CB649128A767A27150849D3684DF92159969A15261ABE78A6DCC549AA027216545F862959B45A22AEA7C01D64630CE498E788BF93C8A2FF4B3DA52387F5B171AF979001917716C2B887FBB0D7D243DBA5337BFA8DA97453236B36FFA278E36A28743A10035B34546514EFFB2CB4CC92F7C614EAE9BCE501B90CD94BB1A1A4C9FAA922BAB2DF8C9BCB8B76CE0B035C8312A990C180F23BD1018D44870021345DB3B39DD55A79229A69B1D79802735E40990CB6092FE614327BDC1FA7B158C16327AAE85456B48D8E64398EF90B24D9C8FBCB9DBC72422126644A83AE4A17069A5CBBCAFBA71473256F1568CB18CC4612A0379BA1E4330A782137282B82ECD25492475F50B30CC7A31521CCB1151CB096944C4E5C21255C1C89C130BE9070DA1696FD1C2901E975B7805C2737BBBAA2965AD10DC43727556A6608E70CD7E59B8FA3CBCD4CB2FE48912B2B72C2D592E0C1AF7702387AF167964685B1D40349F20CA6808B74F87DE67C255145A738F7B3DB8B89F6A71FCF2BB7ED11C6EB2154A37B4CF7063F8BA12433B457F4E3CC4B05370C0A76896A7EE39A863AE49945FAA497766DE3201E34AB7678663E33F85B12C9B570184E0DA3092598AC5B961D7D48C9B30624F829AB66F63B8D34478D35664FB84BECD1951375183B6A2357587F8513C6C5431C0198316AF2CA3987B9F5E219B36116108C8C8E767FE7C690AF248CE0839C13EC4A4CD800FC719BC191B59C5A17EDDBA9CB830474EC0111408243A97335494DDDFA8D5C6A344B0C9ED1061A7D4A840891C24E140BE6D7BDEE01571A1312D4000591F2224B850105878FD3966F537AC348D042B95785B8171F18647619072780C131BDAA462D34A039014012BBB2197A4328F89CC5B750CAD9495BE607D5017DFA105D87C830432795B42C5E89D8B536C087592829ED934822768A1824696A278952627CA2EA73CDA4A306F972B8521B6A7B28DD0AC44E6A8C46BA7291D48574F8356065AFBB073DF9B84B520673F6644D3959BAD576497E0BC846F560CE697C35493A995597F66315639A7B20B84105EAA48F3B4593A0BF12D64AD4CB043B6435D44A2F5E136F21621B88704DF8855C4F1689761A239DFA6874455D48251EB087865AC3825C275A1018B112F01E82D8698F684C8EE127D9D25FB92B180677A36930793DAB0306AAC8BC80A79885B0737C6B319600BB859966AB9F082490E31B669008CD055921CB2360667180611485E5D70E774401739B16F54601D4A600608119E4B3112AAC92D1C094E40318D4259044D16D3AD62B35170B54D16F0F850C593B2062CC3A769B4D28B6A776F981D4BC835612349292B748E30098D29C668A80F68718BF40B7E785CE41C094A7326F3E6992D99B80E64350C5395823B1315160582EC10B01A4337C52C355E98B857B2DF0006A0E8AB49F5BA23B8C7E329C20F91386852A8948751C0F301465D713E7857A7BD8B6076C059BE378F74BB257750F544A31874B79ABEA78F8647E77001506919AA467019E924B7BE53705A5C9411A8B513C8C51300849136E77919CF83CBA9EA1CD6004AA7178C8E7D062E5E13544E060B34984F61C2EB6B56580C45F67DA8982A97BD5BBA9D98462161A55357079C15A2A33992D4B66093DB03676D0C385DCC6DD3C18EC8138738AA0FD6058C9F721C5E3131D7828C9675048949D0CE32E92A3363EC0527F54C89B5140590822D9C96CB1BAA43A8CACCC085EA6FB602B871128AB2C8A2A1644623A1D7175CF4A23F4B497B7BB4D0B652EA4272D7A99A9A5F4C54EAA6482D0BFB4E46EE482B892895C78DCBAE4414543E0BF8AC4C82889BC6DDC385624AC7805CCC644BB430588AA19723AEE5999EF2AE2565C3BB3519E879596C07C24B516C80882627308F4244B98164FDA55AC298B3955EB4C9A35956CA9712BCA0ED46C82CF3506460EB702105E3A4DF7410C9FFEA032F26C02",
        .c = "387F23816CA3AAE900520559080374E6BFF44117D50D78EC4CB4A5C08A955DBBDFCBDE0803301BA3E40A9C33CAAE6DFA71EED311D9ECB84B497D3259B614DA30C71803288904D0E1B9EA8FB5750878602148E6FCE794B91E2778B07D87706E620B727162290CA7CE27DDEC5110770614175CE61A837CE46DA95D6AE7FBD9D4709636807EB67CC5992A643EABA2C253050C7B63CA6CA83617BCAF453AC3E24B9C3617038AA3A2FDC52A32EE7B2EC2B67E5F442BFB8D5429F079096FA7D5B9FBD81B096A75F2C9FAAD8E97DDC003EB88F55F482775BA7BD661BFCFC27D3A7C1071C5A7B020E80F6950BF315A040F4CF025A74B42DC008F4C084690F43E8502BFB740538BF3DF05336DDD6E33A3A4679A33C8A724125A8EC897A593741CA9C934CEEC2F20C82CCF439011C5066937377183C8CC96046985B92F6D014E0223EA542A8BA0A789478EB4B4F04AA98F7E6A4D4AC8852E8670B5A4097F7584FA136A159DAF5005D9FEF08868BFBFBE15D320AEAFFB3DE18C8F39A22A876D750AAD4F85546A88B0D0B498535B590B191894F43763BD13379BBD2CDEF4C98A571103E9409B2227315AB892990984FD6C9F375164959E3035B16C650CE71AFB62F7E3BFC624FACC973B104B44F6A2BCEEE1AACF0C3D8CE9ADDCEB26E3E786C116B70A3CF2C63524F472A2F15DBAB3FF8FDE34F015DD8DCFCBB17793A552E805AD5541ADFC6890B5E3A4DDBBC1CDB9B36D12DC49CA22CBB91F5BB599EF86DC4B5820B8E3D965839DFAA9392D29A878315BEA9A6885994C0D07843A4DC2FB87D0C8507CB17C549C3E60564D263221635F831C9F1F6539F07248E5BA8571541CF69151630161C0DBA57DA3A3C6296FEB02C19E2080F57C657CA74656B88B708450DA5DED44594B12614054D643FEB674E18FD3B863DE526EFB6D2426A58AE363658E8318578029DA3388D5A3C81E53E22F2802CBFD52760AF128EDAE7D1FC3379642FD8B9A791B289D7ECFC546158FBDCFBE1643B03B6874F3BA918849D93E5C0172E684C303D8D58CC5DE703A35FDB5796B81CD415B35DAEB60FCC043827401542FE35AD72375",
        .k = "9A4E2BCC0E5E40528B4C2879D8BE4BF7AA266C16762480054B475814789EABCA",
        .m = "381C547CE06B687A7A17C1A99D11A7A4C35B82685AF51CE452099C60136FD84A",
        .algId = CRYPT_PKEY_ML_KEM,
        .type = CRYPT_KEM_TYPE_MLKEM_512,
    }
};

static const char *MLKEM_SEED_VECTOR = NULL;

static int32_t GetPkey(void *libCtx, const char *attrName, const CMVP_MlkemVector *vector, CRYPT_EAL_PkeyCtx **pkeyPrv,
    CRYPT_EAL_PkeyCtx **pkeyPub)
{
    int32_t ret = CRYPT_CMVP_ERR_ALGO_SELFTEST;
    CRYPT_EAL_PkeyPrv prvKey = { 0 };
    CRYPT_EAL_PkeyPub pubKey = { 0 };
    *pkeyPrv = CRYPT_EAL_ProviderPkeyNewCtx(libCtx, vector->algId, 0, attrName);
    GOTO_ERR_IF_TRUE(*pkeyPrv == NULL, CRYPT_CMVP_ERR_ALGO_SELFTEST);
    *pkeyPub = CRYPT_EAL_ProviderPkeyNewCtx(libCtx, vector->algId, 0, attrName);
    GOTO_ERR_IF_TRUE(*pkeyPub == NULL, CRYPT_CMVP_ERR_ALGO_SELFTEST);

    ret = CRYPT_EAL_PkeySetParaById(*pkeyPrv, vector->type);
    GOTO_ERR_IF_TRUE(ret != CRYPT_SUCCESS, ret);
    ret = CRYPT_EAL_PkeySetParaById(*pkeyPub, vector->type);
    GOTO_ERR_IF_TRUE(ret != CRYPT_SUCCESS, ret);

    prvKey.id = vector->algId;
    prvKey.key.mldsaPrv.data = CMVP_StringsToBins(vector->dk, &(prvKey.key.kemDk.len));
    ret = CRYPT_EAL_PkeySetPrv(*pkeyPrv, &prvKey);
    GOTO_ERR_IF_TRUE(ret != CRYPT_SUCCESS, ret);

    pubKey.id = vector->algId;
    pubKey.key.mldsaPub.data = CMVP_StringsToBins(vector->ek, &(pubKey.key.kemEk.len));
    ret = CRYPT_EAL_PkeySetPub(*pkeyPub, &pubKey);
ERR:
    BSL_SAL_Free(prvKey.key.kemDk.data);
    BSL_SAL_Free(pubKey.key.kemEk.data);
    return ret;
}

static int32_t TestVectorRandom(uint8_t *r, uint32_t rLen)
{
    uint8_t *rand = NULL;
    uint32_t randLen;

    rand = CMVP_StringsToBins(MLKEM_SEED_VECTOR, &randLen);
    if (rand == NULL) {
        return CRYPT_MEM_ALLOC_FAIL;
    }
    if (randLen < rLen) {
        BSL_SAL_Free(rand);
        return CRYPT_CMVP_ERR_ALGO_SELFTEST;
    }

    for (uint32_t i = 0; i < randLen; i++) {
        r[i] = rand[i];
    }
    BSL_SAL_Free(rand);
    return 0;
}

static bool TestMlkemEncapsDecaps(void *libCtx, const char *attrName, const CMVP_MlkemVector *vector)
{
    bool ret = false;
    uint8_t *cipher = NULL;
    uint8_t *cipherVec = NULL;
    uint32_t cipherLen = 0;
    uint32_t cipherVecLen = 0;
    CRYPT_EAL_PkeyCtx *pkeyPrv = NULL;
    CRYPT_EAL_PkeyCtx *pkeyPub = NULL;
    uint8_t sharedKey[32] = { 0 };
    uint32_t sharedLen = sizeof(sharedKey);
    uint8_t *sharedKeyVec = NULL;
    uint32_t sharedKeyVecLen = 0;
    CRYPT_EAL_RandFunc func = CRYPT_RandRegistGet();
    CRYPT_EAL_RandFuncEx funcEx = CRYPT_RandRegistExGet();
    CRYPT_RandRegistEx(NULL);

    GOTO_ERR_IF_TRUE(
        GetPkey(libCtx, attrName, vector, &pkeyPrv, &pkeyPub) != CRYPT_SUCCESS, CRYPT_CMVP_ERR_ALGO_SELFTEST);
    GOTO_ERR_IF_TRUE(CRYPT_EAL_PkeyCtrl(pkeyPrv, CRYPT_CTRL_GET_CIPHERTEXT_LEN, &cipherLen, sizeof(cipherLen)) !=
        CRYPT_SUCCESS, CRYPT_CMVP_ERR_ALGO_SELFTEST);
    cipher = BSL_SAL_Malloc(cipherLen);
    GOTO_ERR_IF_TRUE(cipher == NULL, CRYPT_MEM_ALLOC_FAIL);

    cipherVec = CMVP_StringsToBins(vector->c, &cipherVecLen);
    GOTO_ERR_IF_TRUE(cipherVec == NULL, CRYPT_CMVP_COMMON_ERR);

    sharedKeyVec = CMVP_StringsToBins(vector->k, &sharedKeyVecLen);
    GOTO_ERR_IF_TRUE(sharedKeyVec == NULL, CRYPT_CMVP_COMMON_ERR);

    // regist rand function
    MLKEM_SEED_VECTOR = vector->m;
    CRYPT_RandRegist(TestVectorRandom);
    // encaps
    GOTO_ERR_IF_TRUE(CRYPT_EAL_PkeyEncaps(pkeyPub, cipher, &cipherLen, sharedKey, &sharedLen) != CRYPT_SUCCESS,
        CRYPT_CMVP_ERR_ALGO_SELFTEST);
    GOTO_ERR_IF_TRUE(cipherLen != cipherVecLen, CRYPT_CMVP_ERR_ALGO_SELFTEST);
    GOTO_ERR_IF_TRUE(memcmp(cipher, cipherVec, cipherLen) != 0, CRYPT_CMVP_ERR_ALGO_SELFTEST);
    GOTO_ERR_IF_TRUE(sharedLen != sharedKeyVecLen, CRYPT_CMVP_ERR_ALGO_SELFTEST);
    GOTO_ERR_IF_TRUE(memcmp(sharedKey, sharedKeyVec, sharedLen) != 0, CRYPT_CMVP_ERR_ALGO_SELFTEST);
    // decaps
    GOTO_ERR_IF_TRUE(CRYPT_EAL_PkeyDecaps(pkeyPrv, cipherVec, cipherVecLen, sharedKey, &sharedLen) != CRYPT_SUCCESS,
        CRYPT_CMVP_ERR_ALGO_SELFTEST);
    GOTO_ERR_IF_TRUE(sharedLen != sharedKeyVecLen, CRYPT_CMVP_ERR_ALGO_SELFTEST);
    GOTO_ERR_IF_TRUE(memcmp(sharedKey, sharedKeyVec, sharedLen) != 0, CRYPT_CMVP_ERR_ALGO_SELFTEST);
    ret = true;
ERR:
    BSL_SAL_Free(cipher);
    BSL_SAL_Free(cipherVec);
    BSL_SAL_Free(sharedKeyVec);
    CRYPT_EAL_PkeyFreeCtx(pkeyPrv);
    CRYPT_EAL_PkeyFreeCtx(pkeyPub);
    CRYPT_RandRegist(func);
    CRYPT_RandRegistEx(funcEx);
    return ret;
}

bool CRYPT_CMVP_SelftestMlkemEncapsDecaps(void)
{
    bool ret = TestMlkemEncapsDecaps(NULL, NULL, &MLKEM_VECTOR[0]);
    return ret;
}

bool CRYPT_CMVP_SelftestProviderMlkemEncapsDecaps(void *libCtx, const char *attrName)
{
    return TestMlkemEncapsDecaps(libCtx, attrName, &MLKEM_VECTOR[0]);
}

#endif /* HITLS_CRYPTO_CMVP_ISO19790 || HITLS_CRYPTO_CMVP_FIPS */